import os
import json
from service.SampleService import SampleService
from service.SimilarityService import SimilarityService
from basicTools.FileProcess import FileProcess


class HGMSimController:
    def __init__(self):
        with open("config.json", "r") as file:
            self.config = json.load(file)
        self.sampleSaveDirectory = self.config["storage"]["sampleSaveDirectory"]
        if not os.path.exists(self.sampleSaveDirectory):
            os.makedirs(self.sampleSaveDirectory)

        self.md5SampleSaveDirectory = self.config["storage"]["md5SampleSaveDirectory"]
        if not os.path.exists(self.md5SampleSaveDirectory):
            os.makedirs(self.md5SampleSaveDirectory)

        self.greyImageDirectory = self.config["storage"]["greyImageDirectory"]
        if not os.path.exists(self.greyImageDirectory):
            os.makedirs(self.greyImageDirectory)

        self.sampleService = SampleService()
        self.similarityService = SimilarityService()

    def addSample(self, file, userId, userName):
        savePath = os.path.join(self.sampleSaveDirectory, file.filename)
        if self.sampleService.saveOriginSample(file, savePath):
            fileProcess = FileProcess()
            md5Value = fileProcess.moveAndRenameByMd5(
                self.sampleSaveDirectory, self.md5SampleSaveDirectory, file.filename
            )

            # 生成灰度图
            fileProcess.generateGreyImage(
                self.md5SampleSaveDirectory, self.greyImageDirectory, md5Value
            )

            # generate tlsh value
            tlshvalue = fileProcess.generateTLSH(savePath)

            # 保存数据
            insertId = self.sampleService.saveSample2DB(
                file.filename, userId, userName, md5Value
            )
            insertId1 = self.sampleService.savaFeature2DB(md5Value, None, tlshvalue, 0)

            # 提取样本的图结构
            if insertId != -1 and insertId1 != -1:
                if self.featureInfExtract(
                    md5Value,
                    self.md5SampleSaveDirectory,
                    insertId,
                ):
                    return True

        return False

    def addBatchSample(self, tars, labelInf, userId, userName, isTrainSample):
        """
        批量增加样本，用于数据导入，需要先把样本导入到保存的文件夹下
        """
        for tarName in tars:
            try:
                label = labelInf[tarName]
            except Exception as e:
                label = "SINGLETON"
            if label.startswith("-	[]"):
                label = "SINGLETON"
            fileProcess = FileProcess()
            md5Value = fileProcess.moveAndRenameByMd5(
                self.sampleSaveDirectory, self.md5SampleSaveDirectory, tarName
            )

            # 生成灰度图
            fileProcess.generateGreyImage(
                self.md5SampleSaveDirectory, self.greyImageDirectory, md5Value
            )

            # generate tlsh value
            savePath = os.path.join(self.sampleSaveDirectory, tarName)
            tlshvalue = fileProcess.generateTLSH(savePath)

            # 保存数据
            insertId = self.sampleService.saveSample2DB(
                tarName, userId, userName, md5Value
            )
            insertId1 = self.sampleService.savaFeature2DB(
                md5Value, label, tlshvalue, isTrainSample
            )

            if insertId == -1 or insertId1 == -1:
                print("{}导入失败".format(tarName))

        print("导入成功")

    def queryGreyImage(self, name):
        return self.sampleService.queryGreyImage(name)

    def querySampleInfoListByUser(self, userId):
        """
        请求用户上传的样本信息
        """
        return self.sampleService.querySampleListByUser(userId)

    def queryDashBoardInf(self):
        """
        请求首页信息数据
        """
        res = {}
        res["sampleSum"] = self.sampleService.querySampleSum()[0][0]
        res["famSum"] = self.sampleService.queryfamSum()[0][0]
        res["todayUploadSum"] = self.sampleService.queryTodayUploadSampleSum()[0][0]
        res["labelTop10"] = self.sampleService.queryFamLabelTop10()
        return res

    def featureInfExtract(self, md5Value, savePath, sampleId):
        """
        将样本转换成对应的md5命名
        提取样本特征相关的信息
        todo 增加其余特征的提取
        """
        samplePath = os.path.join(savePath, md5Value)
        if self.sampleService.extractFeature(samplePath, sampleId, md5Value) != -1:
            return True
        return False

    def querySampleInfo(self, sampleId):
        """
        请求样本信息
        """
        pass

    def delSample(self, sampleId):
        """
        删除样本
        """
        pass

    def calculateSamplePairSim(self, sampleId1, sampleId2):
        """
        获得样本对之间的相似性度量结果
        传入的时样本名，后台使用的是md5值
        """
        tar = self.sampleService.querySampleByFileName(sampleId1)
        tar1 = self.sampleService.querySampleByFileName(sampleId2)
        sampleId1Md5 = tar[0][6]
        sampleId2Md5 = tar1[0][6]
        res = {}
        res["simScore"] = self.similarityService.pairSim(sampleId1Md5, sampleId2Md5)
        basicInfs = self.querySamplePairBasicInf(tar, tar1)
        for basicInf in basicInfs.keys():
            res[basicInf] = basicInfs[basicInf]
        print(res)
        return res

    def calculateSampleSimTopk(self, sampleId):
        """
        获得数据库样本中和该样本最相似的topk样本
        """
        tar = self.sampleService.querySampleByFileName(sampleId)
        sampleIdMd5 = tar[0][6]

        kValue = self.config["indicator"]["topK"]
        returnKValue = self.config["indicator"]["returnTopK"]
        ans = self.similarityService.simTopK(sampleIdMd5, kValue, returnKValue)
        res = {}
        res["simScore"] = [item[1] for item in ans]
        basicInfs = []
        for item in ans:
            tar1 = self.sampleService.querySampleByMd5(item[0])
            basicInfs.append(self.querySamplePairBasicInf(tar, tar1))
        res["basicInfs"] = basicInfs
        return res

    def querySamplePairBasicInf(self, tar, tar1):
        """
        获得样本对之间的基础信息
        """
        sampleId1Md5 = tar[0][6]
        sampleId2Md5 = tar1[0][6]
        res = {}
        ans = self.sampleService.calculateReuseFunction(sampleId1Md5, sampleId2Md5)
        for name in ans:
            res[name] = ans[name]

        # [(5, '172709de6edadb750ed668a52b5affcc', 'megasearch', None, 0)]
        famLabel1 = self.sampleService.queryFeatureByMd5(sampleId1Md5)[0][2]
        if famLabel1.startswith("SINGLETON"):
            famLabel1 = "None"
        res["famLabel1"] = famLabel1
        famLabel2 = self.sampleService.queryFeatureByMd5(sampleId2Md5)[0][2]
        if famLabel2.startswith("SINGLETON"):
            famLabel2 = "None"
        res["famLabel2"] = famLabel2
        return res

    def querySimHistoryList(self, operatorName):
        pass

    def querySamplePairInf(self, sampleId1, sampleId2):
        pass


if __name__ == "__main__":
    hGMSimController = HGMSimController()
    # res = hGMSimController.calculateSamplePairSim(
    #     "02d63ae96c49f0bed65dfc64404e2c59", "6e56ab5d8d53a5c8666007530f692f40"
    # )
    # print(res)
    hGMSimController.queryDashBoardInf()
